from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from llm import llm, llm0_1

prompt_gen = '编写一个100字小故事'

prompt = ChatPromptTemplate.from_messages([
    ("system", "你是一个很出色的故事作家."),
    ("user", "{input}")
])

chain_gen = prompt | llm | StrOutputParser()
tale = chain_gen.invoke(prompt_gen)

print(tale)

prompt_gen_q = f'''
根据故事内容，生成一个一句话的简短事实性问题，故事内容如下： ```
{tale}
```
'''
q1 = llm0_1.invoke(prompt_gen_q).content
print('\nQ:', q1)


prompt_qa = f'''
根据故事内容，回答问题，使用一句话的简短回答，故事内容如下： ```
{tale}
```
'''
prompt_qa = ChatPromptTemplate.from_messages([
    ("system", prompt_qa),
    ("user", "{input}")
])

chain_qa = prompt_qa | llm0_1 | StrOutputParser()
ans = chain_qa.invoke(q1)
print('A:', ans)
